# +----------------------------------------------------------------------
# | ChatWork智能聊天办公系统
# +----------------------------------------------------------------------
# | 软件声明: 本系统并非自由软件,未经授权任何形式的商业使用均属非法。
# | 版权保护: 任何企业和个人不允许对程序代码以任何形式任何目的复制/分发。
# | 授权要求: 如有商业使用需求,请务必先与版权所有者取得联系并获得正式授权。
# +----------------------------------------------------------------------
# | Author: ChatWork Team <2474369941@qq.com>
# +----------------------------------------------------------------------
import os
import json
import time
import base64
import logging
import aiofiles
from events import scheduler
from common.cache.queue_cache import QueueCache
from common.enums.apply import DrawsEnum
from common.enums.wallet import WalletEnum
from common.enums.public import CrontabEnum
from common.models.sys import SysCrontabModel
from common.models.draws import DrawsRecordsModel
from common.chain.draws_server import DrawsDalleServer
from common.chain.draws_server import DrawsDoubaoServer
from common.models.users import UserWalletModel
from common.utils.cache import RedisUtil
from common.utils.times import TimeUtil
from common.utils.tools import ToolsUtil
from common.utils.urls import UrlUtil
from plugins.storage.driver import StorageDriver


logger = logging.getLogger(__name__)


async def execute(**kwargs):
    """ 绘画回调任务 """
    start_time = time.time()
    crontab_id: int = int(kwargs["w_id"])
    crontab_job: str = str(kwargs["w_job"])
    if not scheduler.get_job(crontab_job):
        scheduler.pause_job(crontab_job)

    try:
        record_id = await QueueCache.pop_queue_value(QueueCache.DRAW_JOB)
        if record_id is None:
            return True

        await dispense(int(record_id))

        await SysCrontabModel.compute(crontab_id, start_time)
    except Exception as e:
        await SysCrontabModel.compute(crontab_id, start_time, status=CrontabEnum.CRON_ERROR, error=str(e))


async def dispense(record_id: int):
    # 查询任务
    record = await DrawsRecordsModel.filter(id=record_id).first()
    if not record and record.status != DrawsEnum.STATUS_ING:
        return

    # 处理响应
    if record.engine == "doubao":
        await handle_doubao_check(record)
    elif record.engine == "dalle":
        await handle_dalle_check(record)
    elif record.engine == "sd":
        await handle_sd_check()
    else:
        record.status = DrawsEnum.STATUS_FAIL
        record.fail_reason = f"No processing channel {record.engine}"
        await record.save()


async def handle_doubao_check(record):
    """
    豆包回调处理: https://www.volcengine.com/docs/6791/1366783
    """
    album_dir_path = ""
    thumb_dir_path = ""
    try:
        # 限流拦截
        lock_key = f"draws:prevent_{record.id}"
        lock_cache = await RedisUtil.get(lock_key)
        if lock_cache:
            await QueueCache.push_draw(record.id)
            return False

        # 限流锁定
        await RedisUtil.set(lock_key, record.id, 60)
        print(f"豆包绘画『开始』: [重试={record.retry_count}, 任务={record.id}, 时间={time.time()}]")

        # 获取生成结果
        server = DrawsDoubaoServer()
        response = await server.fetch(record.model, record.task_id)

        # 根据状态处理
        if response.get("status") == "done":
            image_url = response.get("image_urls")[0]
            await DrawsRecordsModel.filter(id=record.id).update(
                status=DrawsEnum.STATUS_YES,
                image_url=image_url,
                thumbnail=image_url,
                notify_snap=json.dumps(response, ensure_ascii=False),
                update_time=int(time.time())
            )

            # 图片扩展
            urlName: str = os.path.basename(image_url.split("?")[0])
            fileExt: str = str(urlName.split(".")[-1]).strip().lower()
            if fileExt not in ["png", "jpg", "jpeg", "gif", "webp"]:
                fileExt = "jpg"

            # 存储路径
            filename: str = ToolsUtil.make_md5_str(urlName+str(time.time())) + "." + fileExt
            curr_now: str = TimeUtil.timestamp_to_date(int(time.time()), str("%Y%m%d"))
            album_dir: str = f"storage/applies/draws/albums/{curr_now}/{filename}"
            thumb_dir: str = f"storage/applies/draws/thumbnail/{str(curr_now)}/{filename}"
            album_dir_path: str = UrlUtil.to_root_path(path=album_dir)
            thumb_dir_path: str = UrlUtil.to_root_path(path=thumb_dir)

            # 下载图片
            await ToolsUtil.download_file(url=image_url, save_path=album_dir_path)

            # 生成缩略图
            ToolsUtil.generate_thumbnail(input_path=album_dir_path, output_path=thumb_dir_path)

            # 图上传OSS
            engine: str = await UrlUtil.get_storage_engine()
            if engine != str("local"):
                await StorageDriver.push(path=album_dir_path, key=album_dir)
                await StorageDriver.push(path=thumb_dir_path, key=thumb_dir)

            # 更新图片
            await DrawsRecordsModel.filter(id=record.id).update(
                image_loc=album_dir,
                thumbnail=thumb_dir
            )

            # 绘画成功
            task_time = int(time.time()) - int(record.start_time)
            print(f"豆包绘画『成功』: [重试={record.retry_count}, 任务={record.id}, 耗时={task_time}s]")
        elif response.get("status") == "in_queue":
            await QueueCache.push_draw(record.id)
        else:
            await DrawsRecordsModel.filter(id=record.id).update(
                status=DrawsEnum.STATUS_FAIL,
                fail_reason=response.get("message"),
                notify_snap=json.dumps(response, ensure_ascii=False),
                update_time=int(time.time())
            )
            await UserWalletModel.inc(
                scene=WalletEnum.POINTS,
                user_id=record.user_id,
                change_type=WalletEnum.POINTS_INC_DRAWS_FAIL,
                change_amount=record.points,
                project="AI绘画",
                remarks=response.get("message")
            )
    except Exception as e:
        if record.retry_count < 2:
            print(f"豆包绘画『重试』: [重试={record.retry_count}, 任务={record.id}, 错误={str(e)}]")
            record.retry_count += 1
            record.start_time = int(time.time())
            await record.save()
            await QueueCache.push_draw(record.id)
        else:
            logger.error(f"豆包绘画『失败』: [任务={record.id}, 错误={str(e)}]")
            await DrawsRecordsModel.filter(id=record.id).update(
                status=DrawsEnum.STATUS_FAIL,
                fail_reason=str(e),
                notify_snap=json.dumps({}, ensure_ascii=False),
                update_time=int(time.time())
            )
            if record.points > 0:
                await UserWalletModel.inc(
                    scene=WalletEnum.POINTS,
                    user_id=record.user_id,
                    change_type=WalletEnum.POINTS_INC_DRAWS_FAIL,
                    change_amount=record.points,
                    project="AI绘画",
                    remarks=str(e)
                )
    finally:
        engine: str = await UrlUtil.get_storage_engine()
        if engine != "local":
            if album_dir_path and os.path.exists(album_dir_path):
                os.unlink(album_dir_path)
            if thumb_dir_path and os.path.exists(thumb_dir_path):
                os.unlink(thumb_dir_path)


async def handle_dalle_check(record):
    album_dir_path = ""
    thumb_dir_path = ""
    try:
        print(f"DALLE绘画『开始』: [重试={record.retry_count}, 任务={record.id}, 时间={time.time()}]")

        # 初始通道
        dalle_server = DrawsDalleServer({
            "model": record.model,
            "prompt": record.prompt_in,
            "scale": record.scale,
            "other_params": json.loads(record.other_params)
        })

        # 发起绘画
        response = await dalle_server.imagine()
        if not response.get("data"):
            error = response.get("error", {})
            raise Exception(error.get("message", "响应数据异常"))

        # 存储路径
        filename: str = ToolsUtil.make_md5_str(record.prompt_in + str(time.time())) + ".png"
        curr_now: str = TimeUtil.timestamp_to_date(int(time.time()), str("%Y%m%d"))
        album_dir: str = f"storage/drawing/albums/{curr_now}/{filename}"
        thumb_dir: str = f"storage/drawing/thumbnail/{str(curr_now)}/{str(filename)}"
        album_dir_path: str = UrlUtil.to_root_path(path=album_dir)
        thumb_dir_path: str = UrlUtil.to_root_path(path=thumb_dir)
        if not os.path.isdir(os.path.dirname(album_dir_path)):
            os.makedirs(os.path.dirname(album_dir_path), exist_ok=True)

        # 原图保存
        b64_json = response["data"][0]["b64_json"]
        b64_data = base64.b64decode(b64_json.split(",")[1])
        async with aiofiles.open(album_dir_path, "wb") as file:
            await file.write(b64_data)

        # 生成缩略图
        ToolsUtil.generate_thumbnail(input_path=album_dir_path, output_path=thumb_dir_path)

        # 图上传OSS
        engine: str = await UrlUtil.get_storage_engine()
        if engine != str("local"):
            await StorageDriver.push(path=album_dir_path, key=album_dir)
            await StorageDriver.push(path=thumb_dir_path, key=thumb_dir)

        # 更新状态
        await DrawsRecordsModel.filter(id=record.id).update(
            status=DrawsEnum.STATUS_YES,
            image_loc=album_dir,
            image_url=album_dir,
            thumbnail=thumb_dir,
            notify_snap=json.dumps({}, ensure_ascii=False),
            update_time=int(time.time())
        )

        task_time = int(time.time()) - int(record.start_time)
        print(f"DALLE绘画『成功』: [重试={record.retry_count}, 任务={record.id}, 耗时={task_time}s]")
    except Exception as e:
        if record.retry_count < 2:
            print(f"DALLE绘画『重试』: [重试={record.retry_count+1}, 任务={record.id}, 错误={str(e)}]")
            record.retry_count += 1
            record.start_time = int(time.time())
            await record.save()
            await QueueCache.push_draw(record.id)
        else:
            logger.error(f"DALLE绘画『失败』: [任务={record.id}, 错误={str(e)}]")
            await DrawsRecordsModel.filter(id=record.id).update(
                status=DrawsEnum.STATUS_FAIL,
                fail_reason=str(e),
                notify_snap=json.dumps({}, ensure_ascii=False),
                update_time=int(time.time())
            )
            if record.points > 0:
                await UserWalletModel.inc(
                    scene=WalletEnum.POINTS,
                    user_id=record.user_id,
                    change_type=WalletEnum.POINTS_INC_DRAWS_FAIL,
                    change_amount=record.points,
                    project=str("AI绘画"),
                    remarks=str(e)
                )
    finally:
        engine: str = await UrlUtil.get_storage_engine()
        if engine != "local":
            if album_dir_path and os.path.exists(album_dir_path):
                os.unlink(album_dir_path)
            if thumb_dir_path and os.path.exists(thumb_dir_path):
                os.unlink(thumb_dir_path)


async def handle_sd_check():
    pass
